package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

// ServerConfig 存储服务器连接信息
type ServerConfig struct {
	Host     string
	Port     int
	User     string
	Password string
}

// hostKeyCallback 自动接受所有主机密钥，不处理 known_hosts 文件
func hostKeyCallback(hostname string, remote net.Addr, key ssh.PublicKey) error {
	return nil
}

// preAcceptHostKeys 预先接受所有主机密钥
func preAcceptHostKeys(servers []ServerConfig) error {
	var wg sync.WaitGroup
	errChan := make(chan error, len(servers))
	sem := make(chan struct{}, 10) // 限制并发数

	for _, server := range servers {
		wg.Add(1)
		sem <- struct{}{}
		go func(s ServerConfig) {
			defer func() {
				<-sem
				wg.Done()
			}()

			addr := fmt.Sprintf("%s:%d", s.Host, s.Port)
			config := &ssh.ClientConfig{
				User: s.User,
				Auth: []ssh.AuthMethod{
					ssh.Password(s.Password),
				},
				HostKeyCallback: hostKeyCallback,
				Timeout:         5 * time.Second,
			}

			client, err := ssh.Dial("tcp", addr, config)
			if err != nil {
				errChan <- fmt.Errorf("预连接 %s 失败: %v", addr, err)
				return
			}
			client.Close()
			log.Printf("成功预接受 %s 的主机密钥", addr)
		}(server)
	}

	wg.Wait()
	close(errChan)

	for err := range errChan {
		if err != nil {
			return err
		}
	}
	return nil
}

// readPublicKey 读取公钥文件内容
func readPublicKey(keyPath string) (ssh.PublicKey, error) {
	keyPath = strings.ReplaceAll(keyPath, "~", os.Getenv("HOME"))
	keyData, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("读取公钥文件失败: %w", err)
	}

	pubKey, _, _, _, err := ssh.ParseAuthorizedKey(keyData)
	if err != nil {
		return nil, fmt.Errorf("解析公钥失败: %w", err)
	}

	return pubKey, nil
}

// readServerList 读取服务器列表文件
func readServerList(serverListPath, username string) ([]ServerConfig, error) {
	serverListPath = strings.ReplaceAll(serverListPath, "~", os.Getenv("HOME"))
	data, err := os.ReadFile(serverListPath)
	if err != nil {
		return nil, fmt.Errorf("读取服务器列表文件失败: %w", err)
	}

	lines := strings.Split(string(data), "\n")
	var servers []ServerConfig
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		parts := strings.Split(line, ":")
		if len(parts) != 3 {
			return nil, fmt.Errorf("服务器配置格式错误: %s，期望格式 host:port:password", line)
		}

		var port int
		_, err := fmt.Sscanf(parts[1], "%d", &port)
		if err != nil {
			return nil, fmt.Errorf("解析端口号失败: %s，错误信息: %w", parts[1], err)
		}

		server := ServerConfig{
			Host:     parts[0],
			Port:     port,
			User:     username,
			Password: parts[2],
		}
		servers = append(servers, server)
	}
	return servers, nil
}

// distributeKey 将公钥分发到目标服务器
func distributeKey(ctx context.Context, server ServerConfig, publicKey ssh.PublicKey) error {
	addr := fmt.Sprintf("%s:%d", server.Host, server.Port)
	config := &ssh.ClientConfig{
		User: server.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(server.Password),
		},
		HostKeyCallback: hostKeyCallback,
		Timeout:         5 * time.Second,
	}

	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return fmt.Errorf("连接服务器失败: %w", err)
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return fmt.Errorf("创建会话失败: %w", err)
	}
	defer session.Close()

	// 将公钥追加到服务器的 authorized_keys 文件
	cmd := fmt.Sprintf("mkdir -p ~/.ssh && chmod 700 ~/.ssh && echo '%s' >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys", string(ssh.MarshalAuthorizedKey(publicKey)))
	output, err := session.CombinedOutput(cmd)
	if err != nil {
		return fmt.Errorf("执行命令失败: %w, 输出: %s", err, string(output))
	}

	return nil
}

func main() {
	// 定义命令行参数
	keyPath := flag.String("key", "~/.ssh/id_rsa.pub", "公钥文件路径")
	serverList := flag.String("servers", "servers.txt", "服务器列表文件路径")
	username := flag.String("user", "root", "SSH用户名")
	concurrency := flag.Int("concurrency", 5, "并发连接数")
	timeout := flag.Int("timeout", 10, "连接超时时间(秒)")
	flag.Parse()

	// 验证参数
	if *concurrency < 1 {
		log.Fatal("并发连接数必须大于0")
	}

	// 输出当前工作目录
	workingDir, err := os.Getwd()
	if err != nil {
		log.Fatalf("获取当前工作目录失败: %v", err)
	}
	log.Printf("当前工作目录: %s", workingDir)

	// 输出实际读取的服务器列表文件路径
	log.Printf("正在读取服务器列表文件: %s", *serverList)
	if _, err := os.Stat(*serverList); os.IsNotExist(err) {
		log.Fatalf("服务器列表文件 %s 不存在", *serverList)
	}

	// 读取公钥
	publicKey, err := readPublicKey(*keyPath)
	if err != nil {
		log.Fatalf("读取公钥失败: %v", err)
	}

	// 读取服务器列表
	servers, err := readServerList(*serverList, *username)
	if err != nil {
		log.Fatalf("读取服务器列表失败: %v", err)
	}

	if len(servers) == 0 {
		log.Fatalf("服务器列表文件 %s 内容为空或解析失败", *serverList)
	}

	// 预接受所有主机密钥
	log.Println("正在预接受所有服务器的主机密钥...")
	if err := preAcceptHostKeys(servers); err != nil {
		log.Fatalf("预接受主机密钥失败: %v", err)
	}

	// 创建并发控制结构
	var wg sync.WaitGroup
	serverChan := make(chan ServerConfig, len(servers))
	errChan := make(chan error, *concurrency)
	successCount := 0
	failCount := 0

	// 启动工作协程
	for i := 0; i < *concurrency; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			for server := range serverChan {
				ctx, cancel := context.WithTimeout(context.Background(), time.Duration(*timeout)*time.Second)
				defer cancel()

				err := distributeKey(ctx, server, publicKey)
				if err != nil {
					log.Printf("[Worker %d] 服务器 %s:%d 分发失败: %v", id, server.Host, server.Port, err)
					errChan <- err
				} else {
					log.Printf("[Worker %d] 服务器 %s:%d 分发成功", id, server.Host, server.Port)
					successCount++
				}
			}
		}(i)
	}

	// 向通道发送服务器配置
	for _, server := range servers {
		serverChan <- server
	}
	close(serverChan)

	// 等待所有工作协程完成
	wg.Wait()
	close(errChan)

	// 统计失败次数
	for range errChan {
		failCount++
	}

	log.Printf("密钥分发完成: 成功 %d 台, 失败 %d 台", successCount, failCount)
}
